package main

import (
	"fmt"
	"sync"
	"time"
)

/*
ref: https://studygolang.com/articles/16528
轻松入门流水线FAN模式
*/
func producer(n int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for i := 0; i < n; i++ {
			out <- n
		}

	}()
	return out
}
func square(inCh <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range inCh {
			// simulate
			//time.Sleep(1 * time.Second)
			out <- n * n
		}
	}()
	return out
}
func merge(cs ...<-chan int) <-chan int {
	out := make(chan int) // 适当增加缓冲可提升性能

	var wg sync.WaitGroup
	collect := func(in <-chan int) {
		defer wg.Done()
		for n := range in {
			out <- n
		}
	}
	wg.Add(len(cs))
	// FAN-IN
	for _, c := range cs {
		go collect(c)
	}
	// 错误方式：直接等待是bug，死锁，因为merge写了out，main却没有读
	// wg.Wait()
	// close(out)

	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func main() {
	in := producer(10000)

	startTime := time.Now()

	//3个 square 协程并发运行
	// FAN-OUT
	c1 := square(in)
	c2 := square(in)
	c3 := square(in)

	// consumer
	for ret := range merge(c1, c2, c3) {
		fmt.Printf("%3d ", ret)
	}
	fmt.Println()
	fmt.Println(time.Since(startTime))
}
